import paddle.nn as nn
class BasicBlock(nn.Layer):
    def __init__(self,in_channels,out_channels):
        super(BasicBlock, self).__init__()
        self.block = nn.Sequential(
            nn.Conv2D(in_channels=in_channels, out_channels=out_channels, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU()
            )
        self.short = nn.Sequential(
            nn.Conv2D(in_channels=in_channels,out_channels=out_channels,kernel_size=1,stride=1),
            nn.MaxPool2D(kernel_size=2,stride=2)
        )
    def forward(self,x):
        inputs = x
        x = self.block(x)
        return x + self.short(inputs)
class FCDiscriminator(nn.Layer):
    def __init__(self,num_classes,ndf = 64):
        super(FCDiscriminator,self).__init__()
        self.stage_block = []
        out_channels = [ndf,ndf*2,ndf*4,ndf*8]
        self.block_stage = []
        in_channels = num_classes
        for i in range(4):
            block = BasicBlock(in_channels=in_channels,out_channels=out_channels[i])
            in_channels = out_channels[i]
            self.block_stage.append(block)
        self.classifier = nn.Conv2D(ndf * 8, 1, kernel_size=4, stride=2, padding=1)
        self.up_sample = nn.Upsample(scale_factor=32, mode='bilinear')
    def forward(self, x):
        for stage in self.block_stage:
            x = stage(x)
        x = self.classifier(x)
        x = self.up_sample(x)
        return x